import numpy as np
import matplotlib.pyplot as plt

# 设置matplotlib中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False  # 正常显示负号
from scipy.signal import find_peaks
import pywt  # 小波去噪库（需安装：pip install pywavelets）

# --------------------------
# 1. 生成模拟的甲烷吸收峰信号
# --------------------------
# 波长范围：1652.8 nm ~ 1653.2 nm（覆盖吸收峰）
wavelength = np.linspace(1652.8, 1653.2, 1000)
center_wl = 1653.0  # 吸收峰中心波长
fwhm = 0.08  # 半高宽（模拟值）

# 洛伦兹线型吸收峰（基频信号）
absorbance = 0.8 / (1 + ((wavelength - center_wl) / (fwhm/2))**2)

# 加入激光调制（10 kHz正弦调制），生成二次谐波信号
mod_freq = 10  # 调制频率（kHz）    
mod_depth = 0.02  # 调制深度
harmonic2 = absorbance * np.cos(2 * np.pi * mod_freq * (wavelength - center_wl))**2

# 加入高斯噪声（模拟检测噪声）
np.random.seed(42)  # 固定随机种子，保证结果可复现
noise = np.random.normal(0, 0.05, len(wavelength))  # 噪声幅值0.05
noisy_harmonic2 = harmonic2 + noise

# --------------------------
# 2. 小波去噪优化信号
# --------------------------
# 选用db4小波，4层分解
wavelet = 'db4'
level = 4
coeffs = pywt.wavedec(noisy_harmonic2, wavelet, level=level)

# 硬阈值去噪（保留信号系数，抑制噪声系数）
threshold = 0.1  # 阈值可根据噪声强度调整
coeffs_thresholded = [
    pywt.threshold(c, threshold, mode='hard') if i > 0 else c
    for i, c in enumerate(coeffs)
]

# 重构去噪后的二次谐波信号
denoised_harmonic2 = pywt.waverec(coeffs_thresholded, wavelet)

# --------------------------
# 3. 峰值识别
# --------------------------
# 寻找峰值（设置最小高度和距离，避免误判）
peaks, _ = find_peaks(
    denoised_harmonic2,
    height=0.3,  # 峰值最小高度（过滤小噪声峰）
    distance=50  # 峰值间最小距离（避免同一峰被多次识别）
)

# 获取峰值对应的波长和幅值
peak_wl = wavelength[peaks]
peak_amp = denoised_harmonic2[peaks]

# --------------------------
# 4. 结果可视化
# --------------------------
plt.figure(figsize=(12, 8))

# 子图1：原始吸收峰与带噪声的二次谐波
plt.subplot(2, 2, 1)
plt.plot(wavelength, absorbance, 'b-', label='原始吸收峰（基频）')
plt.title('甲烷吸收峰（洛伦兹线型）')
plt.xlabel('波长 (nm)')
plt.ylabel('吸收率')
plt.legend()

# 子图2：带噪声的二次谐波信号
plt.subplot(2, 2, 2)
plt.plot(wavelength, noisy_harmonic2, 'r-', label='带噪声的二次谐波')
plt.title('含噪声的2f信号')
plt.xlabel('波长 (nm)')
plt.ylabel('信号幅值')
plt.legend()

# 子图3：去噪后的二次谐波信号
plt.subplot(2, 2, 3)
plt.plot(wavelength, denoised_harmonic2, 'g-', label='去噪后的二次谐波')
plt.plot(peak_wl, peak_amp, 'ro', label=f'识别峰值: {peak_wl[0]:.2f} nm')
plt.title('小波去噪后的2f信号与峰值')
plt.xlabel('波长 (nm)')
plt.ylabel('信号幅值')
plt.legend()

# 子图4：峰值放大图
plt.subplot(2, 2, 4)
plt.plot(wavelength, denoised_harmonic2, 'g-')
plt.xlim(1652.9, 1653.1)  # 放大峰值区域
plt.axvline(x=peak_wl[0], color='r', linestyle='--', label=f'峰值波长: {peak_wl[0]:.4f} nm')
plt.title('峰值区域放大')
plt.xlabel('波长 (nm)')
plt.ylabel('信号幅值')
plt.legend()

plt.tight_layout()
plt.savefig('methane_absorption_peak_result.png', dpi=300, bbox_inches='tight')
plt.show(block=False)
plt.pause(22)  # 显示2秒后自动关闭

# 输出识别结果
print(f"识别到的吸收峰峰值波长：{peak_wl[0]:.4f} nm")
print(f"峰值信号幅值：{peak_amp[0]:.4f}")